import UIKit
import Metal
import MetalKit

class MetalCombinedImageProcessor {
    let device: MTLDevice
    private let commandQueue: MTLCommandQueue
    private let extractPipelineState: MTLRenderPipelineState  // 用于提取原图和通道图
    private let alphaPipelineState: MTLRenderPipelineState   // 用于应用透明度
    private let vertexBuffer: MTLBuffer
    private var myDrawable: CAMetalDrawable?
    
    // 定义合并图中通道图的位置和提取方式
    enum MaskPosition {
        case topHalf    // 通道图在上半部分
        case bottomHalf // 通道图在下半部分
        case leftHalf   // 通道图在左半部分
        case rightHalf  // 通道图在右半部分
    }
    
    init() {
        // 获取Metal设备
        guard let device = MTLCreateSystemDefaultDevice() else {
            fatalError("Metal is not supported on this device")
        }
        self.device = device
        
        // 创建命令队列
        guard let commandQueue = device.makeCommandQueue() else {
            fatalError("Failed to create command queue")
        }
        self.commandQueue = commandQueue
        
        // 加载并编译Metal着色器
//        let library = try! device.makeLibrary(source: MetalCombinedShaders.source, options: nil)
        let library = device.makeDefaultLibrary()!
        
        // 获取着色器函数
        let vertexFunction = library.makeFunction(name: "vertexShader")
        let extractFragmentFunction = library.makeFunction(name: "extractFragmentShader")
        let alphaFragmentFunction = library.makeFunction(name: "alphaFragmentShader")
        
        // 创建顶点描述符，定义顶点数据布局
        let vertexDescriptor = MTLVertexDescriptor()

        // 配置第一个属性（position：float4）
        vertexDescriptor.attributes[0].format = .float4    // 数据格式
        vertexDescriptor.attributes[0].offset = 0           // 在顶点结构体中的偏移量
        vertexDescriptor.attributes[0].bufferIndex = 0      // 关联的缓冲区索引

        // 配置第二个属性（texCoord：float2）
        vertexDescriptor.attributes[1].format = .float2
        vertexDescriptor.attributes[1].offset = MemoryLayout<Float>.stride * 4  // 紧跟在position之后
        vertexDescriptor.attributes[1].bufferIndex = 0

        // 配置缓冲区布局（单个顶点的总大小）
        vertexDescriptor.layouts[0].stride = MemoryLayout<Float>.stride * 6
        vertexDescriptor.layouts[0].stepFunction = .perVertex  // 每个顶点对应一个数据
        vertexDescriptor.layouts[0].stepRate = 1
        
        
        // 创建提取用的渲染管道
        let extractPipelineDescriptor = MTLRenderPipelineDescriptor()
        extractPipelineDescriptor.vertexFunction = vertexFunction
        extractPipelineDescriptor.fragmentFunction = extractFragmentFunction
        extractPipelineDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm
        
        // 关键：设置顶点描述符（解决stage_in错误的核心）
        extractPipelineDescriptor.vertexDescriptor = vertexDescriptor
        
        extractPipelineState = try! device.makeRenderPipelineState(descriptor: extractPipelineDescriptor)
        
        // 创建顶点描述符，定义顶点数据布局
        let vertexDescriptor2 = MTLVertexDescriptor()

        // 配置第一个属性（position：float4）
        vertexDescriptor2.attributes[0].format = .float4    // 数据格式
        vertexDescriptor2.attributes[0].offset = 0           // 在顶点结构体中的偏移量
        vertexDescriptor2.attributes[0].bufferIndex = 0      // 关联的缓冲区索引

        // 配置第二个属性（texCoord：float2）
        vertexDescriptor2.attributes[1].format = .float2
        vertexDescriptor2.attributes[1].offset = MemoryLayout<Float>.stride * 4  // 紧跟在position之后
        vertexDescriptor2.attributes[1].bufferIndex = 0

        // 配置缓冲区布局（单个顶点的总大小）
        vertexDescriptor2.layouts[0].stride = MemoryLayout<Float>.stride * 6
        vertexDescriptor2.layouts[0].stepFunction = .perVertex  // 每个顶点对应一个数据
        vertexDescriptor2.layouts[0].stepRate = 1
        
        // 创建透明度处理用的渲染管道
        let alphaPipelineDescriptor = MTLRenderPipelineDescriptor()
        alphaPipelineDescriptor.vertexFunction = vertexFunction
        alphaPipelineDescriptor.fragmentFunction = alphaFragmentFunction
        alphaPipelineDescriptor.colorAttachments[0].pixelFormat = .bgra8Unorm
        
        // 关键：设置顶点描述符（解决stage_in错误的核心）
        alphaPipelineDescriptor.vertexDescriptor = vertexDescriptor2
        // 关键：开启透明混合（Alpha通道生效的核心配置）
        alphaPipelineDescriptor.colorAttachments[0].isBlendingEnabled = true
        alphaPipelineDescriptor.colorAttachments[0].rgbBlendOperation = .add // RGB混合方式：相加
        alphaPipelineDescriptor.colorAttachments[0].alphaBlendOperation = .add // Alpha混合方式：相加
        // 混合因子：符合常规透明逻辑（源像素Alpha控制透明度）
        alphaPipelineDescriptor.colorAttachments[0].sourceRGBBlendFactor = .sourceAlpha
        alphaPipelineDescriptor.colorAttachments[0].sourceAlphaBlendFactor = .sourceAlpha
        alphaPipelineDescriptor.colorAttachments[0].destinationRGBBlendFactor = .oneMinusSourceAlpha
        alphaPipelineDescriptor.colorAttachments[0].destinationAlphaBlendFactor = .oneMinusSourceAlpha
        
        alphaPipelineState = try! device.makeRenderPipelineState(descriptor: alphaPipelineDescriptor)
        
        // 设置顶点数据（全屏矩形）
        let vertices: [Float] = [
            // 位置(x,y,z,w)      // 纹理坐标(u,v)
            -1.0, -1.0, 0.0, 1.0,  0.0, 1.0,
             1.0, -1.0, 0.0, 1.0,  1.0, 1.0,
            -1.0,  1.0, 0.0, 1.0,  0.0, 0.0,
             1.0,  1.0, 0.0, 1.0,  1.0, 0.0
        ]
        vertexBuffer = device.makeBuffer(bytes: vertices, length: vertices.count * MemoryLayout<Float>.stride, options: [])!
    }
    
    // 从合并图中处理并生成原图大小的透明图
    func process(combinedTexture: MTLTexture, maskPosition: MaskPosition, drawable: CAMetalDrawable) {
        
        myDrawable = drawable
         
        // 提取原图和通道图
        guard let originalTexture = extractOriginalImage(from: combinedTexture, position: maskPosition),
              let maskTexture = extractMaskImage(from: combinedTexture, position: maskPosition) else {
            print("提取原图或通道图失败")
            return
        }
        
        // 将通道图缩放到原图尺寸
        guard let scaledMaskTexture = scaleTexture(maskTexture, to: originalTexture.size, device: device) else {
            print("缩放通道图失败")
            return
        }
        
        // 应用透明度
        applyAlpha(originalTexture: originalTexture, maskTexture: scaledMaskTexture)
        

    }
    
    // 从合并图中提取原图
    private func extractOriginalImage(from combinedTexture: MTLTexture, position: MaskPosition) -> MTLTexture? {
        let originalSize = getOriginalImageSize(combinedSize: CGSize(width: combinedTexture.width, height: combinedTexture.height), position: position)
        
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .bgra8Unorm,
            width: Int(originalSize.width),
            height: Int(originalSize.height),
            mipmapped: false
        )
        textureDescriptor.usage = [.renderTarget, .shaderRead]
        
        guard let originalTexture = device.makeTexture(descriptor: textureDescriptor) else {
            return nil
        }
        
        // 创建提取区域的纹理坐标
        let (minTexCoord, maxTexCoord) = getOriginalImageTexCoords(position: position)
        
        // 执行提取操作
        renderToTexture(
            sourceTexture: combinedTexture,
            destinationTexture: originalTexture,
            pipelineState: extractPipelineState,
            texCoordAdjustment: (minTexCoord, maxTexCoord)
        )
        
        return originalTexture
    }
    
    // 从合并图中提取通道图
    private func extractMaskImage(from combinedTexture: MTLTexture, position: MaskPosition) -> MTLTexture? {
        let maskSize = getMaskImageSize(combinedSize: CGSize(width: combinedTexture.width, height: combinedTexture.height), position: position)
        
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .bgra8Unorm,
            width: Int(maskSize.width),
            height: Int(maskSize.height),
            mipmapped: false
        )
        textureDescriptor.usage = [.renderTarget, .shaderRead]
        
        guard let maskTexture = device.makeTexture(descriptor: textureDescriptor) else {
            return nil
        }
        
        // 创建提取区域的纹理坐标
        let (minTexCoord, maxTexCoord) = getMaskImageTexCoords(position: position)
        
        // 执行提取操作
        renderToTexture(
            sourceTexture: combinedTexture,
            destinationTexture: maskTexture,
            pipelineState: extractPipelineState,
            texCoordAdjustment: (minTexCoord, maxTexCoord)
        )
        
        return maskTexture
    }
    
    // 将纹理缩放到指定尺寸
    private func scaleTexture(_ texture: MTLTexture, to size: CGSize, device: MTLDevice) -> MTLTexture? {
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .bgra8Unorm,
            width: Int(size.width),
            height: Int(size.height),
            mipmapped: false
        )
        textureDescriptor.usage = [.renderTarget, .shaderRead]
        
        guard let scaledTexture = device.makeTexture(descriptor: textureDescriptor) else {
            return nil
        }
        
        // 使用线性滤波进行缩放
        renderToTexture(
            sourceTexture: texture,
            destinationTexture: scaledTexture,
            pipelineState: extractPipelineState,
            texCoordAdjustment: (float2(0, 0), float2(1, 1))
        )
        
        return scaledTexture
    }
    
    // 应用透明度到原图
    private func applyAlpha(originalTexture: MTLTexture, maskTexture: MTLTexture) {
        
        // 创建命令缓冲区
        guard let commandBuffer = commandQueue.makeCommandBuffer() else {
            fatalError("无法创建命令缓冲区")
        }
        
        // 创建渲染通行证描述符
        let renderPassDescriptor = MTLRenderPassDescriptor()
        renderPassDescriptor.colorAttachments[0].texture = myDrawable?.texture // 绑定目标纹理
        renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 0) // 清除色
        renderPassDescriptor.colorAttachments[0].loadAction = .clear // 渲染前清除为 clearColor
        renderPassDescriptor.colorAttachments[0].storeAction = .store // 渲染后保存结果
        
        // 创建渲染命令编码器
        guard let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor) else {
            fatalError("无法创建渲染命令编码器")
        }
        
        // 设置渲染管道状态
        renderEncoder.setRenderPipelineState(alphaPipelineState)
        
        // 设置顶点缓冲区
        renderEncoder.setVertexBuffer(vertexBuffer, offset: 0, index: 0)
        
        // 设置纹理
        renderEncoder.setFragmentTexture(originalTexture, index: 0)
        renderEncoder.setFragmentTexture(maskTexture, index: 1)
        
        // 绘制矩形
        renderEncoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
        
        // 结束编码
        renderEncoder.endEncoding()
        
        // 提交命令并等待完成
        commandBuffer.present(myDrawable!)
        commandBuffer.commit()
        
        
    }
    
    // 通用渲染到纹理的方法
    private func renderToTexture(
        sourceTexture: MTLTexture,
        destinationTexture: MTLTexture,
        pipelineState: MTLRenderPipelineState,
        texCoordAdjustment: (min: float2, max: float2)
    ) {
        // 创建命令缓冲区
        guard let commandBuffer = commandQueue.makeCommandBuffer() else {
            fatalError("无法创建命令缓冲区")
        }
        
        // 创建渲染通行证描述符
        let renderPassDescriptor = MTLRenderPassDescriptor()
        renderPassDescriptor.colorAttachments[0].texture = destinationTexture
        renderPassDescriptor.colorAttachments[0].loadAction = .clear
        renderPassDescriptor.colorAttachments[0].clearColor = MTLClearColor(red: 0, green: 1, blue: 0, alpha: 0)
        renderPassDescriptor.colorAttachments[0].storeAction = .store
        
        // 创建渲染命令编码器
        guard let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor) else {
            fatalError("无法创建渲染命令编码器")
        }
        
        // 设置渲染管道状态
        renderEncoder.setRenderPipelineState(pipelineState)
        
        // 设置顶点缓冲区
        renderEncoder.setVertexBuffer(vertexBuffer, offset: 0, index: 0)
        
        // 设置纹理坐标调整参数
        var adjustment = texCoordAdjustment
        renderEncoder.setFragmentBytes(&adjustment, length: MemoryLayout<(float2, float2)>.stride, index: 0)
        
        // 设置源纹理
        renderEncoder.setFragmentTexture(sourceTexture, index: 0)
        
        // 绘制矩形
        renderEncoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
        
        // 结束编码
        renderEncoder.endEncoding()
        
        // 提交命令并等待完成
//        commandBuffer.present(destinationTexture)
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()
    }
    
    // 获取原图尺寸
    private func getOriginalImageSize(combinedSize: CGSize, position: MaskPosition) -> CGSize {
//        switch position {
//        case .topHalf, .bottomHalf:
//            return CGSize(width: combinedSize.width, height: combinedSize.height / 2)
//        case .leftHalf, .rightHalf:
//            return CGSize(width: combinedSize.width / 2, height: combinedSize.height)
//        }
            
        return CGSize(width: combinedSize.width * 0.66, height: combinedSize.height)
    }
    
    // 获取通道图尺寸
    private func getMaskImageSize(combinedSize: CGSize, position: MaskPosition) -> CGSize {
        // 通道图可以是任意尺寸，这里直接返回其在合并图中的实际尺寸
//        switch position {
//        case .topHalf, .bottomHalf:
//            return CGSize(width: combinedSize.width, height: combinedSize.height / 2)
//        case .leftHalf, .rightHalf:
//            return CGSize(width: combinedSize.width / 2, height: combinedSize.height)
//        }
        return CGSize(width: combinedSize.width * 0.33, height: combinedSize.height * 0.5)
    }
    
    // 获取原图的纹理坐标范围
    private func getOriginalImageTexCoords(position: MaskPosition) -> (float2, float2) {
        switch position {
        case .topHalf:      // 原图在下半部分
            return (float2(0, 0.5), float2(1, 1))
        case .bottomHalf:   // 原图在上半部分
            return (float2(0, 0), float2(1, 0.5))
        case .leftHalf:     // 原图在右半部分
            return (float2(0.5, 0), float2(1, 1))
        case .rightHalf:    // 原图在左半部分
            return (float2(0, 0), float2(0.66, 1))
        }
    }
    
    // 获取通道图的纹理坐标范围
    private func getMaskImageTexCoords(position: MaskPosition) -> (float2, float2) {
        switch position {
        case .topHalf:      // 通道图在上半部分
            return (float2(0, 0), float2(1, 0.5))
        case .bottomHalf:   // 通道图在下半部分
            return (float2(0, 0.5), float2(1, 1))
        case .leftHalf:     // 通道图在左半部分
            return (float2(0, 0), float2(0.5, 1))
        case .rightHalf:    // 通道图在右半部分
            return (float2(0.66, 0), float2(1, 0.5))
        }
    }
}

// 辅助结构体
struct Vertex {
    let position: float4
    let texCoord: float2
}

